# 实现Test1
import time
from dolfinx import mesh, plot, fem, default_real_type, geometry
from dolfinx.fem import functionspace, petsc
from mpi4py import MPI
import pyvista as pv
import numpy as np
import basix
from basix.ufl import element, mixed_element
import ufl
import norm
from petsc4py import PETSc

import tools  # 定义了一些坐标计算相关的函数


def marker(x): return np.isclose(x[0], 1) & np.isclose(x[1], 1)


def true_velocity_function(x):
    """
    速度真解 函数
    """
    u1 = -4 * x[1] * (4 - x[0]**2)**2 * (4 - x[1]**2)
    u2 = 4 * x[0] * (4 - x[0]**2) * (4 - x[1]**2)**2
    return [u1, u2]


def true_velocity_function_ufl(x):
    """
    速度真解 ufl函数
    """
    return ufl.as_vector(true_velocity_function(x))


def true_pressure_function(x):
    """
    压强真解 函数
    """
    return 150 * np.sin(x[0])


def true_pressure_function_ufl(x):
    """
    压强真解 函数
    """
    return 150 * ufl.sin(x[0])


def lambda_function(x):
    """拉格朗日乘子真解

    Args:
        x (_type_): msh
    """
    return [np.exp(x[0]), np.exp(x[1])]


def lambda_function_ufl(x):
    """拉格朗日乘子真解
        ufl形式
    Args:
        x (_type_): msh
    """
    return ufl.as_vector([ufl.exp(x[0]), ufl.exp(x[1])])


def lambda_x_dx(x, y):
    return np.exp(x)

def lambda_x_dy(x, y):
    return np.zeros_like(y)

def lambda_y_dx(x, y):
    return np.zeros_like(x)

def lambda_y_dy(x, y):
    return np.exp(y)




def laplace_u_ufl(x):
    """拉普拉斯算子作用于速度真解
    """
    u1 = 8 * x[1] * (3 * x[0]**4 + 6 * x[0]**2 * x[1]
                     ** 2 - 48*x[0]**2 - 8*x[1]**2 + 80)
    u2 = -8 * x[0] * (6 * x[0]**2 * x[1]**2 - 8 * x[0] **
                      2 + 3 * x[1]**4 - 48 * x[1]**2 + 80)
    return ufl.as_vector([u1, u2])


def laplace_u(x):
    # 提取x和y分量（注意：x[0]是x坐标，x[1]是y坐标）
    x_coord = x[0]
    y_coord = x[1]
    
    # 计算u1的拉普拉斯算子（对应SymPy的Laplacian of u1）
    term1 = 24 * y_coord * (x_coord**2 - 4)**2
    term2 = 16 * y_coord * (3 * x_coord**2 - 4) * (y_coord**2 - 4)
    u1 = term1 + term2
    
    # 计算u2的拉普拉斯算子（对应SymPy的Laplacian of u2）
    term3 = -16 * x_coord * (x_coord**2 - 4) * (3 * y_coord**2 - 4)
    term4 = -24 * x_coord * (y_coord**2 - 4)**2
    u2 = term3 + term4
    
    return np.array([u1, u2])


def grad_p(coords):
    """压力梯度真解 - 动态维度版本
    """
    gp1 = 150 * np.cos(coords[:, 0])
    gp2 = np.zeros(coords.shape[0])
    return np.array([gp1, gp2])


def grad_p_ufl(x):
    """压力梯度真解
    """
    return ufl.as_vector([150 * ufl.cos(x[0]), 0])  # 压强真解的梯度，y方向为0


def displacement_function_ufl(x):
    """位移真解
    """
    return true_velocity_function_ufl(x)  # 位移真解和速度真解相同


def displacement_function(x):
    """位移真解
    """
    return true_velocity_function(x)  # 位移真解和速度真解相同


def c(mu, Y):
    """双线性形式 c(μ, Y) 这里ufl.sym()的使用存在困惑 是否有必要?
    """
    A = (ufl.inner(ufl.grad(mu), ufl.grad(Y))) * dx_solid
    B = ufl.inner(mu, Y) * dx_solid
    return A + B


def d(x):
    # 位移约束 d = u(X̄) - X
    A = np.array(true_velocity_function(
        identity_map(displacement_function(x))))
    B = np.array(displacement_function(x))
    return A - B


def g(x):
    """固体域体力项 g = -ΔX - λ + g_∂B UFL表达式"""
    neg_laplace_X = -laplace_u_ufl(x)  # UFL支持直接对向量取负
    neg_lambda = -lambda_function_ufl(x)

    # 计算 g_∂B = ∇X · n_solid_vec
    X_expr = displacement_function(x)
    grad_X = ufl.grad(displacement_function_ufl(x))
    # g_boundary = ufl.dot(grad_X, n_solid_vec)

    # 组合各项（UFL向量运算）
    # return neg_laplace_X + neg_lambda + g_boundary
    return neg_laplace_X + neg_lambda


def identity_map(s):
    return s  # s 为固体参考域坐标 (x, y)


def fluid_boundary(x):
    """流体边界"""
    return np.isclose(x[0], 2.0) | np.isclose(
        x[0], -2.0) | np.isclose(x[1], -2.0) | np.isclose(x[1], 2.0)


def solid_boundary(x):
    """固体边界"""
    return np.isclose(x[0], 1.0) | np.isclose(
        x[0], -1.0) | np.isclose(x[1], -1.0) | np.isclose(x[1], 1.0)


def u_x_derivative_x(x, y):
    # 计算u_x(x, y)对x的导数
    return 16 * x * y * (4 - x**2) * (4 - y**2)


def u_x_derivative_y(x, y):
    # 计算u_x(x, y)对y的导数
    return (12 * y ** 2 - 16) * (4 - x ** 2)**2


def u_y_derivative_x(x, y):
    # 计算u_y(x, y)对x的导数
    return (16 - 12*x**2) * (4 - y**2)**2


def u_y_derivative_y(x, y):
    # 计算u_y(x, y)对y的导数
    return -16 * x * y * (4 - x**2) * (4 - y**2)

def p_derivative_x(x, y):
    # 计算p(x, y)对x的导数
    return 150 * np.cos(x)

def p_derivative_y(x, y):
    return np.zeros_like(y)

def g_partial_B(x, y):
    if np.all(np.isclose(x, -1)):
        A1 = -u_x_derivative_x(-1, y)
        A2 = -u_y_derivative_x(-1, y)
    elif np.all(np.isclose(x, 1)):
        A1 = u_x_derivative_x(1, y)
        A2 = u_y_derivative_x(1, y)
    elif np.all(np.isclose(y, -1)):
        A1 = -u_x_derivative_y(x, -1)
        A2 = -u_y_derivative_y(x, -1)
    elif np.all(np.isclose(y, 1)):
        A1 = u_x_derivative_y(x, 1)
        A2 = u_y_derivative_y(x, 1)
    return np.array([A1, A2])


def lambda_partial_B1(x, y):
    return np.array([np.zeros_like(x), -np.exp(y)])

def lambda_partial_B2(x, y):
    return np.array([np.exp(x), np.zeros_like(y)])

def lambda_partial_B3(x, y):
    return np.array([np.zeros_like(x), np.exp(y)])

def lambda_partial_B4(x, y):
    return np.array([-np.exp(x), np.zeros_like(y)])

def lambda_partial_B(x, y):
    if np.all(np.isclose(y, -1)):
        return lambda_partial_B1(x, y)
    elif np.all(np.isclose(x, 1)):
        return lambda_partial_B2(x, y)
    elif np.all(np.isclose(y, 1)):
        return lambda_partial_B3(x, y)
    elif np.all(np.isclose(x, -1)): 
        return lambda_partial_B4(x, y)

def getStiffnessMatrix(A: PETSc.Mat):
    """
    获取刚度矩阵的值(np.array)
    """
    # 需要导入
    # from petsc4py import PETSc
    m, n = A.getSize()
    matrix_data = np.zeros((m, n))

    # 获取矩阵的所有非零元素
    row, col, val = A.getValuesCSR()

    # 将CSR格式的数据转换为密集矩阵
    for i in range(m):
        for j in range(row[i], row[i + 1]):
            matrix_data[i, col[j]] = val[j]

    return matrix_data



X_order = 1  # 固体位移元阶数
# h_solid = 0.5  # 固体网格大小
# h_fluid = 0.5
h_solid = 0.125  # 固体网格大小
h_fluid = 0.125


Lx_solid, Ly_solid = 2, 2  # 固体网格长宽
Lx_fluid, Ly_fluid = 4, 4  # 流体网格长宽
n_solid = [int(Lx_solid / h_solid), int(Ly_solid / h_solid)]
n_fluid = [int(Lx_fluid / h_fluid), int(Ly_fluid / h_fluid)]
num_solid_cells = 2 * n_solid[0]**2  # 2*N_solid*N_solid
num_fluid_cells = 2 * n_fluid[0]**2  # 2*N_fluid*N_fluid
v_size = (2 * n_fluid[0] + 1)**2 * 2  # 流体速度自由度数量
p_size = (n_fluid[0] + 1)**2  # 流体压力自由度数量
lambda_size = 2 * (n_solid[0] + 1)**2  # 拉格朗日乘子自由度数量


if X_order == 1:
    X_size = 2 * (n_solid[0] + 1)**2  # 固体位移自由度数量
    print("Use P2/P1/P1/P1 element.")
elif X_order == 2:
    X_size = 2 * (2 * n_solid[0] + 1)**2
    print("Use P2/P1/P2/P1 element.")


# 打印网格信息

print(f"Solid mesh size: {n_solid}, Number of cells: {num_solid_cells}.")
print(f"Fluid mesh size: {n_fluid}, Number of cells: {num_fluid_cells}.")

solid_msh = mesh.create_rectangle(
    MPI.COMM_WORLD, [[-1, -1], [1, 1]], n=n_solid, cell_type=mesh.CellType.triangle)
fluid_msh = mesh.create_rectangle(
    MPI.COMM_WORLD, [[-2, -2], [2, 2]], n=n_fluid, cell_type=mesh.CellType.triangle)

# key为边界边索引 value为边界边的单元索引
solid_boundary_dict = tools.search_boundary_edge_cell(
    solid_msh, counterclock=True)

# 定义有限元空间 https://defelement.org/elements/p1-iso-p2.html
u_p2 = element(
    "Lagrange", fluid_msh.basix_cell(), degree=2, shape=(fluid_msh.geometry.dim,), dtype=default_real_type
)
p_p1 = element("Lagrange", fluid_msh.basix_cell(),
               degree=1, dtype=default_real_type)
X_p1 = element(
    "Lagrange", solid_msh.basix_cell(), degree=1, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)
X_p2 = element(
    "Lagrange", solid_msh.basix_cell(), degree=2, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)
lambda_p1 = element(
    "Lagrange", solid_msh.basix_cell(), degree=1, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)

P1 = element(
    "Lagrange", solid_msh.basix_cell(), degree=1,  dtype=default_real_type
)  # 定义一个普通的P1元(非向量元) 用于计算基函数值
P2 = element(
    "Lagrange", solid_msh.basix_cell(), degree=2,  dtype=default_real_type
)  # 定义一个普通的P2元(非向量元) 用于计算基函数值

dx_fluid = ufl.Measure("dx", domain=fluid_msh)
dx_solid = ufl.Measure("dx", domain=solid_msh)
V = functionspace(fluid_msh, u_p2)  # 流体速度函数空间
Q = functionspace(fluid_msh, p_p1)  # 流体压力函数空间
if X_order == 1:
    W = functionspace(solid_msh, X_p1)  # 固体位移函数空间
elif X_order == 2:
    W = functionspace(solid_msh, X_p2)  # 固体位移函数空间

S = functionspace(solid_msh, lambda_p1)  # 固体拉格朗日乘子函数空间


# 流体域测试函数和试函数
(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)
(X, lambda_) = ufl.TrialFunction(W), ufl.TrialFunction(S)
(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)
(Y, mu) = ufl.TestFunction(W), ufl.TestFunction(S)
#

u_bc = fem.Function(V)  # 真实速度函数空间
p_bc = fem.Function(Q)  # 真实压力函数空间
u_bc.interpolate(true_velocity_function)
p_bc.interpolate(true_pressure_function)

# 设置流体边界自由度
fdim = fluid_msh.topology.dim - 1
fluid_facets = mesh.locate_entities_boundary(fluid_msh, fdim, fluid_boundary)
solid_facets = mesh.locate_entities_boundary(solid_msh, fdim, solid_boundary)


x_fluid = ufl.SpatialCoordinate(fluid_msh)  # 流体域上的空间坐标
x_solid = ufl.SpatialCoordinate(solid_msh)  # 固体域上的空间坐标

# 流体域体力项 f（包含 -Δu + ∇p + c(λ, v)）
# f_fluid = -ufl.div(ufl.grad(u)) + ufl.grad(p) + c(lambda_, v)

# 固体域体力项 g（包含 -ΔX - λ + g_∂B）

# g_boundary = ufl.dot(ufl.grad(X), n_solid_vec)
# g_solid = -ufl.div(ufl.grad(X)) - lambda_ + g_boundary  # g_boundary 需提前定义

# 4 * 4矩阵 a_{i}_{j} i, j = 1, 2, 3, 4
a_1_1 = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx_fluid  # a_f(u, v) a_1_1
a_1_2 = -ufl.inner(ufl.div(v), p) * dx_fluid  # (p, v)项
a_1_3 = None
a_1_4 = None

a_2_1 = -ufl.inner(ufl.div(u), q) * dx_fluid
a_2_2 = fem.Constant(fluid_msh, 0.0) * ufl.inner(p, q) * dx_fluid  # (p, q)项
a_2_3 = None
a_2_4 = None

a_3_1 = None
a_3_2 = None
a_3_3 = ufl.inner(ufl.grad(X), ufl.grad(Y)) * dx_solid  # a_s(X, Y)
a_3_4 = -c(lambda_, Y)  # c(λ, Y)

a_4_1 = None
a_4_2 = None
a_4_3 = -c(mu, X)  # c(μ, X)
a_4_4 = fem.Constant(solid_msh, 0.0) * ufl.inner(lambda_, mu) * dx_solid
# P2/P1/P1/P1 流体32*32 固体16*16剖分 刚度矩阵 10695*10695
a = fem.form(
    [
        [a_1_1, a_1_2, a_1_3, a_1_4],
        [a_2_1, a_2_2, a_2_3, a_2_4],
        [a_3_1, a_3_2, a_3_3, a_3_4],
        [a_4_1, a_4_2, a_4_3, a_4_4],
    ]
)


A = petsc.assemble_matrix_block(a)  # petsc4py.PETSc.Mat
A.assemble()
A_np = getStiffnessMatrix(A)

# 求解前准备数据
num_basis_func_X = X_p1.degree * 3
num_basis_func_lambda = lambda_p1.degree * 3
num_fluid_func_u = u_p2.degree * 3
G = tools.GaussIntegrate(order=9)
bb_tree = geometry.bb_tree(fluid_msh, 2)
bb_tree_s = geometry.bb_tree(solid_msh, 2)

fluid_msh.topology.create_connectivity(2, 0)  # 创建流体域顶点和边的连接
solid_msh.topology.create_connectivity(2, 0)  # 创建固体域顶点和边的连接
lambda_dofmap = S.dofmap  # 拉格朗日乘子的自由度映射 因为计算的是c(lambda, v)
v_dofmap = V.dofmap  # 流体速度的自由度映射
X_dofmap = W.dofmap  # 固体位移的自由度映射
p_dofmap = Q.dofmap  # 流体压力的自由度映射

V_coords = V.tabulate_dof_coordinates()[:, :-1]  # 流体速度自由度坐标(P2元 流体单元)
S_coords = S.tabulate_dof_coordinates()[:, :-1]  # 拉格朗日乘子自由度坐标(P1元 固体单元)
X_coords = W.tabulate_dof_coordinates()[:, :-1]  # 固体位移自由度坐标(P1元 固体单元)
# 拆分变量
lambda_basis_values_x, lambda_basis_values_y = G.get_vector_basis_values(
    num_basis_func=num_basis_func_lambda, element=lambda_p1)
u_basis_values_x, u_basis_values_y = G.get_vector_basis_values(
    num_basis_func=num_fluid_func_u, element=u_p2)
X_basis_values_x, X_basis_values_y = G.get_vector_basis_values(
    num_basis_func=num_basis_func_lambda, element=lambda_p1)

Cf = np.zeros((2*(n_solid[0] + 1)**2, 2*(2*n_fluid[0] + 1)**2))


# 定义载荷向量
# L = [f 0 g d]^T
L1 = np.zeros(v_size)
L2 = np.zeros(p_size)
L3 = np.zeros(X_size)
L4 = np.zeros(lambda_size)

# 组装 C_f
# 尺寸578(mu) * 8450(u) 假设固体单元16*16 流体单元32*32

for i in range(num_solid_cells):
    lambda_idx = lambda_dofmap.cell_dofs(i)  # 获取固体单元的全局自由度索引 P1元 长度为3
    solid_cell_points = S_coords[lambda_idx]  # 获取固体单元的自由度坐标 P1元 对应顶点
    bdt = tools.BasisDerivativeTransformer(solid_cell_points)
    det_J = bdt.detJ

    gs_points_phy = G.to_physical_gauss_points(*solid_cell_points)
    gs_points_ref = bdt.local_to_ref(*gs_points_phy.T)
    P1_basis_value, P1_basis_value_d_xi, P1_basis_value_d_eta = P1.tabulate(
        1, gs_points_ref.T)
    P2_basis_value, P2_basis_value_d_xi, P2_basis_value_d_eta = P2.tabulate(
        1, gs_points_ref.T)
    # 计算对应的导数
    P1_basis_value_d_x = bdt.local_dx1(
        P1_basis_value_d_xi, P1_basis_value_d_eta)
    P1_basis_value_d_y = bdt.local_dy1(
        P1_basis_value_d_xi, P1_basis_value_d_eta)
    P2_basis_value_d_x = bdt.local_dx1(P2_basis_value_d_xi, P2_basis_value_d_eta)
    P2_basis_value_d_y = bdt.local_dy1(P2_basis_value_d_xi, P2_basis_value_d_eta)
    # 计算固体单元高斯点所在的流体单元
    # order=3 三个高斯点都在边中点 会出现共享单元的情况
    # case1 T_f_k k=1; 三个固体高斯点只对应一个流体单元
    fluid_cell_idx_ = tools.search_element(gs_points_phy, fluid_msh, bb_tree)
    fluid_cell_idx = tools.process_case1_search_cell(fluid_cell_idx_)
    u_idx = v_dofmap.cell_dofs(fluid_cell_idx)  # 查询速度网格对应的索引

    # 计算流体单元的高斯积分点对应的基函数值
    for alpha in range(num_fluid_func_u):  # 速度P2元 基函数有6个
        for beta in range(num_basis_func_lambda):  # 拉格朗日乘子P1元 基函数有3个
            # fenicsx向量元是交错排列的 x0, y0, x1, y1, ...
            
            
            # 注释梯度项: 只保留item_1, item_4的第一项 item_2, item_3都设置为0
            item_1 = P1_basis_value[:, beta] * P2_basis_value[:, alpha]
            # item_1 += P1_basis_value_d_x[:, beta] * P2_basis_value_d_x[:, alpha]

            # item_2 = P2_basis_value_d_y[:, alpha] * P1_basis_value_d_y[:, beta]
            # item_3 = P2_basis_value_d_x[:, alpha] * P1_basis_value_d_x[:, beta]
            
            item_4 = P1_basis_value[:, beta] * P2_basis_value[:, alpha]
            # item_4 += P1_basis_value_d_y[:, beta] * P2_basis_value_d_y[:, alpha]

            Cf[2 * lambda_idx[beta], 2 * u_idx[alpha]] += det_J * np.dot(item_1, G.weights)
            Cf[2 * lambda_idx[beta] + 1, 2 * u_idx[alpha] + 1] += det_J * np.dot(item_4, G.weights)

            # # Cf[2 * lambda_idx[beta], 2 * u_idx[alpha]+1
            # #    ] += det_J * np.dot(item_2, G.weights)
            # # Cf[2 * lambda_idx[beta]+1, 2 * u_idx[alpha]
            # #    ] += det_J * np.dot(item_3, G.weights)
            

    # 组装g  载荷向量 这里拉格朗日乘子和固体位移都是同一套网格 都使用P1元
    for beta in range(num_basis_func_X):
        # 这里补偿 (laplace(lambda), Y)
        # g_vec = -laplace_u(gs_points_phy.T) - lambda_function(gs_points_phy.T)  原始版本
        g_vec = -laplace_u(gs_points_phy.T) - lambda_function(gs_points_phy.T)
        L3[2 * lambda_idx[beta]] += det_J * \
            np.dot(g_vec[0] * X_basis_values_x[:, beta], G.weights)
        L3[2 * lambda_idx[beta] + 1] += det_J * \
            np.dot(g_vec[1] * X_basis_values_y[:, beta], G.weights)


# 组装g_partialB 载荷向量L3 注意这里是个1D积分 不是二维积分
G1d = tools.GaussIntegrate1D(order=9)
for bd_edge, edge_info in solid_boundary_dict.items():
    cell_idx = edge_info["cell_index"]
    X_idx = X_dofmap.cell_dofs(cell_idx)
    nodes = X_coords[X_idx][:3]

    bdt = tools.BasisDerivativeTransformer(nodes)

    gs_points_phy = np.zeros((2, G1d.points.size))
    # 变换回标准坐标
    p1, p2 = edge_info["coordinates"]
    if abs(p1[0] - p2[0]) < 1e-10:  # y方向的边界边
        gs_points_phy[0] = p1[0]
        gs_points_phy[1] = G1d.get_nodes(p1[1], p2[1])
        a, b = np.sort([p1[1], p2[1]])
    else:
        gs_points_phy[0] = G1d.get_nodes(p1[0], p2[0])
        gs_points_phy[1] = p1[1]
        a, b = np.sort([p1[0], p2[0]])
    ref_points = bdt.local_to_ref(*gs_points_phy)
    # 高斯点的函数值
    f1, f2 = g_partial_B(*gs_points_phy)  # 用于g的计算
    l1, l2 = lambda_partial_B(*gs_points_phy)
    for beta in range(num_basis_func_X):
        rr_x = (f1 + l1) * P1.tabulate(0, ref_points.T)[0][:, beta]
        rr_y = (f2 + l2) * P1.tabulate(0, ref_points.T)[0][:, beta]
        L3[2 * X_idx[beta]] += G1d.integrate(a, b, rr_x)
        L3[2 * X_idx[beta]+1] += G1d.integrate(a, b, rr_y)


# 组装f 载荷向量 L1
flag = False
for i in range(num_fluid_cells):
    u_idx = v_dofmap.cell_dofs(i)  # 速度索引 长度为6
    fluid_cell_points = V_coords[u_idx]  # 获取流体单元的自由度坐标 P2元
    bdt = tools.BasisDerivativeTransformer(fluid_cell_points[:3])
    det_J = bdt.detJ
    gs_points_phy = G.to_physical_gauss_points(
        *fluid_cell_points[:3])  # 传入三角形坐标
    
    # 计算对应的固体单元idx
    # 使用重心坐标查找
    bary_center = np.mean(fluid_cell_points[:3], axis=0)
    solid_cell_idx = tools.search_element(bary_center, solid_msh, bb_tree_s)
    if solid_cell_idx.array.size != 0:
        flag = True  # 流体区域对应固体区域
        
    for beta in range(num_fluid_func_u):
        f_vec = -laplace_u(gs_points_phy.T) + \
            grad_p(gs_points_phy)  # 2 * n_gs_points
        if flag:
            # 最新怀疑
            f_vec += np.exp(gs_points_phy.T)
        L1[2 * u_idx[beta]] += det_J * \
            np.dot(f_vec[0] * u_basis_values_x[:, beta], G.weights)
        L1[2 * u_idx[beta] + 1] += det_J * \
            np.dot(f_vec[1] * u_basis_values_y[:, beta], G.weights)
    
    flag = False


load_vec = np.concatenate([L1, L2, L3, L4])
# 将结果传回刚度矩阵
A_np = tools.assign_to_corner(A_np, Cf, corner="lower_left")
A_np = tools.assign_to_corner(A_np, Cf.T, corner="upper_right")

# 设置边界条件
v_dofs = fem.locate_dofs_topological(
    V=V, entity_dim=1, entities=fluid_facets)  # 速度的边界节点
p_dofs = fem.locate_dofs_topological(
    V=Q, entity_dim=1, entities=fluid_facets)  # 压力的边界节点
for v_dof in v_dofs:
    A_np[2 * v_dof] = 0
    A_np[2 * v_dof + 1] = 0
    # 这里的数值应该是对应位置的边界值 本例使用零边界条件
    load_vec[2 * v_dof] = 0
    load_vec[2 * v_dof + 1] = 0
    A_np[2 * v_dof, 2 * v_dof] = 1
    A_np[2 * v_dof + 1, 2 * v_dof + 1] = 1


# 固定压力为0
p_idx = 0
p_fix_x, p_fix_y = Q.tabulate_dof_coordinates()[p_idx][:-1]  # 固定压力点的坐标
A_np[v_size+p_idx] = 0
A_np[v_size+p_idx, v_size+p_idx] = 1
load_vec[v_size+p_idx] = true_pressure_function([p_fix_x, p_fix_y])




# solve
# 稀疏求解能差不多快10倍
# x = np.linalg.solve(A_np, load_vec)  # 使用numpy的线性代数求解器
x = tools.sparse_solve(A_np, load_vec)


# 拆分解 u p X lambda
u_sol = x[:v_size]
p_sol = x[v_size:v_size+p_size]
X_sol = x[v_size+p_size:-lambda_size]
lambda_sol = x[-lambda_size:]



# 构造真解
u_true = fem.Function(V)
p_true = fem.Function(Q)
X_true = fem.Function(W)
lambda_true = fem.Function(S)
# 插值
u_true.interpolate(true_velocity_function)
p_true.interpolate(true_pressure_function)
X_true.interpolate(displacement_function)
lambda_true.interpolate(lambda_function)

# 计算最大节点误差                                                                                                             
u_max_err = norm.get_max_node_err(
    u_true_func=true_velocity_function,
    u_num=u_sol,
    func_space=V
)
p_max_err = norm.get_max_node_err(
    u_true_func=true_pressure_function,
    u_num=p_sol,
    func_space=Q
)
X_max_err = norm.get_max_node_err(
    u_true_func=displacement_function,
    u_num=X_sol,
    func_space=W
)
lambda_max_err = norm.get_max_node_err(
    u_true_func=lambda_function,
    u_num=lambda_sol,
    func_space=S
)
print(f"Velocity max error: {u_max_err:.6e}")
print(f"Pressure max error: {p_max_err:.6e}")
print(f"Displacement max error: {X_max_err:.6e}")
print(f"Lagrange multiplier max error: {lambda_max_err:.6e}\n")


# 计算L2误差
u_l2_err = norm.get_L2_err(
    msh=fluid_msh,
    dof_map=v_dofmap,
    num_cells=num_fluid_cells,
    u_true_func=true_velocity_function,
    u_num=u_sol,
    func_space=V,
    order=2
)
p_l2_err = norm.get_L2_err(
    msh=fluid_msh,
    dof_map=p_dofmap,
    num_cells=num_fluid_cells,
    u_true_func=true_pressure_function,
    u_num=p_sol,
    func_space=Q,
    order=1 
)
X_l2_err = norm.get_L2_err(
    msh=solid_msh,
    dof_map=X_dofmap,
    num_cells=num_solid_cells,
    u_true_func=displacement_function,
    u_num=X_sol,
    func_space=W,
    order=X_order
)
lambda_l2_err = norm.get_L2_err(
    msh=solid_msh,
    dof_map=lambda_dofmap,
    num_cells=num_solid_cells,
    u_true_func=lambda_function,
    u_num=lambda_sol,
    func_space=S,
    order=1
)
print(f"Velocity L2 error: {u_l2_err:.6e}")
print(f"Pressure L2 error: {p_l2_err:.6e}")
print(f"Displacement L2 error: {X_l2_err:.6e}")
print(f"Lagrange multiplier L2 error: {lambda_l2_err:.6e}\n")


# H1
u_h1_err = norm.get_H1_err(
    msh=fluid_msh,
    dof_map=v_dofmap,
    num_cells=num_fluid_cells,
    u_true_func_diff_ls=[u_x_derivative_x, u_x_derivative_y, u_y_derivative_x, u_y_derivative_y],
    u_num=u_sol,
    func_space=V,
    order=2
)
p_h1_err = norm.get_H1_err(
    msh=fluid_msh,
    dof_map=p_dofmap,
    num_cells=num_fluid_cells,
    u_true_func_diff_ls=[p_derivative_x, p_derivative_y],
    u_num=p_sol,
    func_space=Q,
    order=1
)
X_h1_err = norm.get_H1_err(
    msh=solid_msh,
    dof_map=X_dofmap,
    num_cells=num_solid_cells,
    u_true_func_diff_ls=[u_x_derivative_x, u_x_derivative_y, u_y_derivative_x, u_y_derivative_y],
    u_num=X_sol,
    func_space=W,
    order=X_order
)
lambda_h1_err = norm.get_H1_err(
    msh=solid_msh,
    dof_map=lambda_dofmap,
    num_cells=num_solid_cells,
    u_true_func_diff_ls=[lambda_x_dx, lambda_x_dy, lambda_y_dx, lambda_y_dy],
    u_num=lambda_sol,
    func_space=S,
    order=1
)
print(f"Velocity H1 error: {u_h1_err:.6e}")
print(f"Pressure H1 error: {p_h1_err:.6e}")
print(f"Displacement H1 error: {X_h1_err:.6e}")
print(f"Lagrange multiplier H1 error: {lambda_h1_err:.6e}\n")
